package motiffinder;

import io.FastaReadWrite;

import java.util.ArrayList;
import java.util.Random;
import java.util.Vector;

public class MotifFinder {

	String [] sequences = null;
	int width = 0;
	Vector<Vector<Double>> motifMatrix = null; // pattern description
	ArrayList<Integer> positions = new ArrayList<Integer>();
	
	char [] patterns = {'A', 'T', 'G', 'C'};
	double bj = 0.01; // residue, arbitrary constant
	int N = 0;

	// tests the program.
	public static void main(String [] args)
	{
		Vector<String> sequence = new Vector<String>();
		Vector<Integer> idx = new Vector<Integer>();
		Vector<String> motifString = new Vector<String>();
		Vector<Integer> motifLen = new Vector<Integer>();
		Vector<Integer> duplicates = new Vector<Integer>();
		System.out.println("Ready to read");
		FastaReadWrite.read(sequence, idx, motifString, motifLen, duplicates);
		System.out.println("Done reading");
		String [] seqArr = new String[10];
		for (int i = 200; i < sequence.size() && i < 210; i++)
		{
			seqArr[i-200] = sequence.get(i);
		}
		
		MotifFinder mf = new MotifFinder();
		mf.setSequences(seqArr, 6); // just test one set.
		System.out.println("Ready to run!");
		mf.run(10000, 500, 2); // every 400 iterations, check windows of size 4, sample one
		System.out.println("Done!");
	}
	/**
	 * width is length of motif, sequences is set of sequence to play with.
	 * @param sequences
	 * @param width
	 */
	public void setSequences(String [] sequences, int width)
	{
		Random random = new Random();
		this.sequences = sequences;
		this.N = sequences.length;
		this.width = width;
		motifMatrix = new Vector<Vector<Double>>();
		for (int r = 0; r < width; r++)
		{
			Vector<Double> currCol = new Vector<Double>();
			for (int c = 0; c < 4; c++)
			{
				currCol.add(0.25); // initialize to a constant
			}
			motifMatrix.add(currCol);
		}
		for (int s = 0; s < sequences.length; s++) 
		{
			int r = (int)( (sequences[0].length() - width) * random.nextDouble() );
			positions.add(r); // random start sequences
		}
	}
	
	/**
	 * M denotes how many times to 'shift'
	 * @param numIter
	 * @param M
	 */
	public void run(int numIter, int M, int window)
	{
		Random r = new Random();
		int count = 1;
		int seen = 0;
		for (int iter = 0; iter < numIter; iter++)
		{
			int seqIdx = (int)(sequences.length * r.nextDouble());
			predictiveUpdate(seqIdx);
			sampling(seqIdx);
			if ((iter + 1) % M == 0)
			{
				seen++;
				if (count == seen)
				{
					phaseShift(window, seqIdx);
					seen = 0;
					count++;
				} // harder for this to be called in later phases.
			}
		}
		System.out.println("The following is motif matrix : ");
		for (int i = 0; i < motifMatrix.size(); i++)
			System.out.println(motifMatrix.get(i));
		System.out.println("And the following is starting positions : " );
		System.out.println(positions);
	}

	/**
	 * checks phase shifts
	 * @param m
	 * @param window
	 */
	private void phaseShift(int window, int seqIdx) 
	{
		Random random = new Random();

		for (int pidx = 0; pidx < positions.size(); pidx++)
		{
			double cum = 0.0;
			Integer pos = positions.get(pidx);
			ArrayList<Double> Akx = new ArrayList<Double>();
			int startIdx = (pos - window > 0) ? pos - window : 0;
			int endIdx = (pos + window < sequences[0].length() - width) ? pos + window : sequences[0].length() - width;
			for (int i = startIdx; i < endIdx; i++)
			{
				double logAx = 0.0;
				this.logSegmentWeight(seqIdx, i, logAx);
				Akx.add(Math.exp(logAx));
				cum += Math.exp(logAx);
			}

			// resample
			double r = cum * random.nextDouble();
			double curr = 0.0;
			for (int i = 0; i < Akx.size(); i++)
			{
				curr += Akx.get(i);
				if (curr > r) 
				{
					positions.set(pidx, startIdx + i);
					break;
				}
			}
		}

	}
	private void sampling(int seqIdx) {
		
		sampleByAx(seqIdx);
		// the following samples on F, formula (2).
//		sampleByF(seqIdx);
		
	}
	private void sampleByF(int seqIdx) 
	{
		Random random = new Random();
		double [] segmentWeight = new double[sequences[seqIdx].length() - width];
		double sampleCum = 0.0;
		double min = 0.0;
		for (int f = 0; f < segmentWeight.length; f++)
		{
			double segCum = 0.0;
			for (int i = 0; i < width; i++)
			{
				int ci [] = new int[4]; // atcg
				for (int p = 0; p < 4; p++)
				{
					for (int s = 0; s < sequences.length; s++)
					{
						int currIdx = f + i; // position of cij.
						if (sequences[s].charAt(currIdx) == patterns[p]) ci[p]++; // increment if same pattern
					}
					double qij = this.motifMatrix.get(i).get(p);
					double currTerm = ci[p] * (Math.log(qij / 0.25));
					segCum += currTerm;
				} //pattern count
			}
			segmentWeight[f] = segCum;
			sampleCum += segCum;
			if (f == 0 || min > segCum)
			{
				min = segCum;
			}
		} // calculate cij * log(qij / pj)
		for (int f = 0; f < segmentWeight.length; f++)
		{
			sampleCum += Math.abs(min + 5.0);
			segmentWeight[f] += Math.abs(min + 5.0);
		} // reorient cum
		double r = sampleCum * random.nextDouble();
		double seenCum = 0.0;
		for (int f = 0; f < segmentWeight.length; f++)
		{
			seenCum += segmentWeight[f];
			if (seenCum > r)
			{
				this.positions.set(seqIdx, f);
				break;
			}
		}
	}
	
	/**
	 * Samples by random pick on probability
	 * @param seqIdx
	 */
	private void sampleByAx(int seqIdx) 
	{
		double [] segmentWeight = new double[sequences[seqIdx].length() - width];
		double cum = 0.0;
		Random random = new Random();
		for (int i = 0; i < segmentWeight.length; i++)
		{
			double logAx = 0; // current segment probability
			logAx = logSegmentWeight(seqIdx, i, logAx);
			double Ax = Math.exp(logAx);
			cum += Ax;
			segmentWeight[i] = Ax;
		} // calculate each segment weight
		
		double currSelection = random.nextDouble() * cum;
		double seenCum = 0.0;
		for (int i = 0; i < segmentWeight.length; i++)
		{
			seenCum += segmentWeight[i];
			if (seenCum > currSelection)
			{
				positions.set(seqIdx, i);
				break;
			}
		}
	}
	private double logSegmentWeight(int seqIdx, int i, double logAx) {
		for (int w = 0; w < width; w++)
		{
			int currSidx = w + i;
			char currPattern = sequences[seqIdx].charAt(currSidx);
			int pidx = -1;
			if (currPattern == patterns[0]) pidx = 0;
			else if (currPattern == patterns[1]) pidx = 1;
			else if (currPattern == patterns[2]) pidx = 2;
			else if (currPattern == patterns[3]) pidx = 3;
			
			logAx += Math.log(motifMatrix.get(w).get(pidx)); // Px is 0.25 as base probability
		}
		return logAx;
	}
	
	
	private void predictiveUpdate(int seqIdx) {
		// TODO Auto-generated method stub
		for (int i = 0; i < width; i++)
		{
			int ci [] = new int[4]; // atcg
			for (int p = 0; p < 4; p++)
			{
				for (int s = 0; s < sequences.length; s++)
				{
					int currIdx = this.positions.get(s) + i; // position of cij.
					if (s == seqIdx)
					{
						
					}
					else
					{
						if (sequences[s].charAt(currIdx) == patterns[p]) ci[p]++; // increment if same pattern
					}
				}
			} //pattern count
			double [] qi = new double[4];
			double motifCum = 0.0;
			for (int p = 0; p < 4; p++)
			{
				qi[p] = (ci[p] + bj) / ((double)(N - 1 + sequences.length * bj));
				motifMatrix.get(i).set(p, qi[p]); // set motif matrix
			}
//			for (int p = 0; p < 4; p++)
//			{
//				qi[p] = (ci[p] + bj) / ((double)(N - 1 + sequences.length * bj));
//				motifMatrix.get(i).set(p, qi[p] / motifCum);
//			} // normalize
		}
	}
	
	public Vector<Vector<Double>> getMotifMatrix()
	{
		return this.motifMatrix;
	}
	
	public ArrayList<Integer> getSites()
	{
		return this.positions;
	}
}
